#include "test.h"
#include <math.h> 
#include <map>

test::test() {}

test::test(double v) {
    test_value.push_back(1);
    test_double_data.push_back(v);
}

size_t test::getTotalSize() const {
    return std::accumulate(test_value.begin(), test_value.end(), (size_t)1,
                         std::multiplies<size_t>());
}

test::test(int dim, size_t shape[], double data[]) {
    std::copy(shape, shape + dim, std::back_inserter(test_value));
    std::copy(data, data + getTotalSize(), std::back_inserter(test_double_data));
}

int test::get_dim() const { 
    return test_value.size(); 
}

size_t *test::get_shape_array() {
    return test_value.empty() ? nullptr : &test_value[0];
}

double *test::get_data_array() { 
    return &test_double_data[0]; 
}

std::vector<double>& test::get_data_vec(){
    return  this->test_double_data;
};

template <class T> 
test helper(const test &lhs, const test &rhs) {
    test result;
    result.test_value = lhs.test_value;
    size_t totalSize = lhs.getTotalSize();

    result.test_double_data.reserve(totalSize);
    std::transform(lhs.test_double_data.begin(), lhs.test_double_data.begin() + totalSize,
                  rhs.test_double_data.begin(), std::back_inserter(result.test_double_data), T());

    return result;
}

test operator+(const test &lhs, const test &rhs) {
    return helper<std::plus<double>>(lhs, rhs);
}

test operator-(const test &lhs, const test &rhs) {
    return helper<std::minus<double>>(lhs, rhs);
}

test operator*(const test &lhs, const test &rhs) {
    test result;
    size_t lhsSize = lhs.getTotalSize();
    size_t rhsSize = rhs.getTotalSize();

    if (lhsSize == 1 || rhsSize == 1) {
      double scalar = (lhsSize == 1 ? lhs : rhs).test_double_data[0];
      const test &t = lhsSize == 1 ? rhs : lhs;
      result.test_value = t.test_value;
      result.test_double_data.reserve(lhsSize * rhsSize);
      std::transform(t.test_double_data.begin(), t.test_double_data.end(),
                    std::back_inserter(result.test_double_data),
                    [&scalar](double x) { return x * scalar; });
    } else {
      size_t I = lhs.test_value[0];
      size_t J = rhs.test_value[1];
      size_t K = lhs.test_value[1];
      result.test_value = {I, J};
      size_t totalSize = result.getTotalSize();
      result.test_double_data.resize(totalSize);

      for (size_t i = 0; i < I; ++i) {
        for (size_t j = 0; j < J; ++j) {
          for (size_t k = 0; k < K; ++k) {
            result.test_double_data[J * i + j] +=
                lhs.test_double_data[K * i + k] * rhs.test_double_data[J * k + j];
          }
        }
      }
    }
    
    return result;
}

std::ostream &operator<<(std::ostream &os, const test &t) {
    std::ostream_iterator<double> it(os, " ");
    std::copy(t.test_value.begin(), t.test_value.end(), it);
    std::copy(t.test_double_data.begin(), t.test_double_data.end(), it);

    return os;
}

test ReLU(const test &x){
    test result;

    result.test_value = x.test_value;
    for(size_t i = 0; i < x.getTotalSize(); i++ ){
        result.test_double_data.push_back(x.test_double_data[i]*(x.test_double_data[i]>0));
    }

    return result;
}

test Flatten(const test &x){
    test result;
    size_t xSize = x.getTotalSize();
//    std::cout << "====Flatten=====" << x.test_value[0] << "----" << x.test_value[1]  << "----" << x.test_value[2]  << "----"  << x.test_value[3]  << "----"  <<std::endl;
    size_t shape_1 = xSize/x.test_value[0];
    std::vector<size_t> shape_;
    shape_.push_back(x.test_value[0]);
    shape_.push_back(shape_1);
    result.test_value = shape_;
    result.test_double_data = x.test_double_data;

    return result;
}

test Inut2d(const test &x){
     test result;

//    std::cout << "====Inut2d=====" << x.test_value[0] << "----" << x.test_value[1]  << "----" << x.test_value[2]  << "----"  << x.test_value[3]  << "----"  <<std::endl;
//   std::cout << "====Inut2d=====" << x.test_double_data[0] << "----" << x.test_double_data[1]  << "----" << x.test_double_data[2]  << "----"  << x.test_double_data[3]  << "----"  <<std::endl;

     std::vector<size_t> shape_;
     shape_.push_back(x.test_value[0]);
     shape_.push_back(x.test_value[3]);
     shape_.push_back(x.test_value[1]);
     shape_.push_back(x.test_value[2]);
     result.test_value = shape_;
     for(size_t i = 0; i < x.test_value[0]; i++ ){
        for(size_t j =0; j < x.test_value[3];j++){
            for(size_t k = 0; k <x.test_value[1]*x.test_value[2]; k++){
                result.test_double_data.push_back(x.test_double_data[j + k*x.test_value[3] + i*x.test_value[3]*x.test_value[1]*x.test_value[2]]);
            }
        }
    }

    return result;
}

test Linear(const test &w, const test &x, const test &b){
    test result_wT;
    std::vector<size_t> shape_0;

//    std::cout << "====x.test_value[0]=====" << x.test_value[0] << "----" << x.test_value[1]  << "----" << x.test_value[2]  << "----"  << x.test_value[3]  << "----"  <<std::endl;
//    std::cout << "====w.test_value[0]=====" << w.test_value[0] << "----" << w.test_value[1]  << "----" << w.test_value[2]  << "----"  << w.test_value[3]  << "----"  <<std::endl;
//    std::cout << "====b.test_value[0]=====" << b.test_value[0] << "----" << b.test_value[1]  << "----" << b.test_value[2]  << "----"  << b.test_value[3]  << "----"  <<std::endl;

    shape_0.push_back(x.test_value[1]);
    shape_0.push_back(w.test_value[0]);
    size_t totalSize_w = w.getTotalSize();
    result_wT.test_double_data.reserve(totalSize_w);
    result_wT.test_value = shape_0;
    for(size_t i = 0; i < w.test_value[1]; i++){
        for(size_t j = 0; j < w.test_value[0]; j++){
            result_wT.test_double_data.push_back(w.test_double_data[i + j*w.test_value[1]]);
        }
    }
    test result_xwT;
    std::vector<size_t> shape_xwT;
    shape_xwT.push_back(x.test_value[0]);
    shape_xwT.push_back(result_wT.test_value[1]);
    size_t totalSize_xwT = x.test_value[0]*result_wT.test_value[1];
    result_xwT.test_double_data.resize(totalSize_xwT);
    result_xwT.test_value = shape_xwT;
    result_xwT =  x * result_wT;

    test result_bBrodcast;
    std::vector<size_t> shape_1;
    shape_1.push_back(b.test_value[0]);
    shape_1.push_back(1*x.test_value[0]);
    result_bBrodcast.test_value = shape_1;
    for(size_t i = 0; i < result_bBrodcast.test_value[0]; i++){
        for(size_t j = 0; j < result_bBrodcast.test_value[1]; j++){
            result_bBrodcast.test_double_data.push_back(b.test_double_data[i]);
        }
    }

   test result_bBroadcastT;
   std::vector<size_t> shape_2;
   size_t totalSize_b = result_bBrodcast.getTotalSize();
   result_bBrodcast.test_double_data.resize(totalSize_b);
   shape_2.push_back(result_bBrodcast.test_value[1]);
   shape_2.push_back(result_bBrodcast.test_value[0]);
   result_bBroadcastT.test_value = shape_2;

    for(size_t i = 0; i < result_bBrodcast.test_value[1]; i++){
        for(size_t j = 0; j < result_bBrodcast.test_value[0]; j++){
            result_bBroadcastT.test_double_data.push_back(result_bBrodcast.test_double_data[i + j*result_bBrodcast.test_value[1]]);
        }
    }

    test result;
    result.test_double_data.resize(result_xwT.test_value[0]*result_xwT.test_value[1]);
    result = result_xwT + result_bBroadcastT;

    return result;
}

std::vector<double> maxpooling(std::vector<std::vector<double> > map, int k_w, int k_h, int s_w, int s_h)
{
    int row = map.size(); int col = map[0].size();
	int out_row = 1+(row-k_h)/s_h, out_col = 1+(col-k_w)/s_w;
	int mod_row = row%s_h, mod_col = col%s_w;
	std::vector<std::vector<double> > res(out_row, std::vector<double>(out_col, 0));
	std::vector<std::vector<double> > pad_map(map);

	if (mod_row != 0) out_row++;
	if (mod_col != 0) out_col++;

	for (int i=0;i<row;i++) {
		for (int j = 0; j < k_h - mod_row; j++)
		{
			pad_map[i].push_back(map[i][col - 1]);
		}

        for (int j = 0; j < k_w - mod_col; j++)
        {
            pad_map.push_back(pad_map[row - 1]);
        }
    }

    for(int i=0;i<out_row;i++) {
		for (int j = 0; j < out_col; j++)
		{
			int start_x = j*s_w;
			int start_y = i*s_h;
			std::vector<double> temp;

			for(int ii=0;ii<k_w;ii++) {
				for (int jj = 0; jj < k_h; jj++)
				{
					temp.push_back(pad_map[start_y + jj][start_x + ii]);
				}
            }
			sort(temp.begin(), temp.end());
			res[i][j] = temp[temp.size() - 1];
		}
    }

    std::vector<double> max_temp;
    for (auto l : res)
    {
        for (auto ll : l)
	    {
		    max_temp.push_back(ll);
	    }
	}

    return max_temp;
}

test MaxPool2d(const test &x, const int kernel_size, const int stride){
 //   std::cout <<"=====sssssssss ======="<<std::endl;
 //   std::cout << x.test_value[0] <<"  "<< x.test_value[1] << "  " << x.test_value[2] << "  "  << x.test_value[3] << "  "  << std::endl;
 //   std::cout << "====" << kernel_size << "  " << std::endl;

    test max2d_temp;
    size_t xSize = x.getTotalSize();
    std::cout << xSize << std::endl;
    std::vector<size_t> temp;
    temp.push_back(x.test_value[0]);
    temp.push_back(x.test_value[1]);
    temp.push_back(floor(x.test_value[2]/kernel_size));
    temp.push_back(floor(x.test_value[3]/kernel_size));
    max2d_temp.test_value = temp;

    for(size_t i = 0; i < x.test_value[0]; i++ ){
      for(size_t j =0; j < x.test_value[1];j++){
        std::vector<std::vector<double> > tmp;
        for(size_t k =0; k < x.test_value[2];k++){
          std::vector<double> tmp_;
          for(size_t l =0; l < x.test_value[3];l++){
            double int_data = x.test_double_data[l + k*x.test_value[3] + j*x.test_value[3]*x.test_value[2] + i*x.test_value[3]*x.test_value[2]*x.test_value[1] ];
            tmp_.push_back(int_data);
          }
          tmp.push_back(tmp_);
        }
       auto res = maxpooling(tmp,kernel_size,kernel_size,stride,stride);
       for(int i = 0; i< (int)res.size(); i++){
            max2d_temp.test_double_data.push_back(res[i]);
        }
      }
    }

    return max2d_temp;
}

test Conv2d(const test &x, const int in_channels, const int out_channels, const int kernel_size, const test &w, const test &b){
    test conv2d_temp;
    std::cout <<"=====Conv2d======="<<std::endl;
    std::cout << "x.test_value[0]=" << x.test_value[0] <<" x.test_value[1] "<< x.test_value[1] << " x.test_value[2] " << x.test_value[2] << " x.test_value[3] "  << x.test_value[3] << "  "  << std::endl;
    std::cout << "w.test_value[0]=" << w.test_value[0] <<" w.test_value[1] "<< w.test_value[1] << " w.test_value[2] " << w.test_value[2] << " w.test_value[3] "  << w.test_value[3] << "  "  << std::endl;
    std::cout << "b.test_value[0]=" << b.test_value[0] <<" b.test_value[1] "<< b.test_value[1] << " b.test_value[2] " << b.test_value[2] << " b.test_value[3] "  << b.test_value[3] << "  "  << std::endl;

    size_t out_w = x.test_value[2]- kernel_size + 1;
    size_t out_h = x.test_value[3]- kernel_size + 1;

    std::vector<size_t> temp;
    temp.push_back(x.test_value[0]);
    temp.push_back(w.test_value[0]);
    temp.push_back(out_w);
    temp.push_back(out_h);
    conv2d_temp.test_value = temp;

    size_t out_map = out_w * out_h;
    size_t kernel_map = kernel_size * kernel_size;

       

    return conv2d_temp;
}